import numpy as np
import scipy.linalg
import os
import sys
from math import pi
from matplotlib.pylab import *

if __name__=='__main__':

    t = np.loadtxt('time.dat');
    rho0 = np.loadtxt("rho.dat")
    H  = np.loadtxt("hamiltonian_real.dat")
    H += np.loadtxt("hamiltonian_imag.dat")
    
    a =  1.0
    c = -1.0
    b = np.sqrt(8.0)

    E,U=scipy.linalg.eig(H)

    print "Hamiltonian ",H
    print "Eigenvalues ",E
    Uh = U.T.conj()

    rho = rho0.copy()
    f = np.zeros((t.shape[0],4),dtype=np.complex)
    i = 0
    f[i,0] = rho[0,0] 
    f[i,1] = rho[1,1] 
    f[i,2] = rho[0,1] 
    f[i,3] = rho[1,0] 
    
    for j in range(1,t.shape[0]):
        dt = t[j]-t[j-1]
        u = scipy.linalg.matfuncs.expm2(-1j*H*dt)
        #u = np.dot(scipy.linalg.inv(np.eye(2,2) + 1j*H*dt/2),np.eye(2,2) - 1j*H*dt/2)

        rho = np.dot(u,np.dot(rho,u.T.conj()))

        i = j
        f[i,0] = rho[0,0] 
        f[i,1] = rho[1,1] 
        f[i,2] = rho[0,1] 
        f[i,3] = rho[1,0] 



    # Solution by the qk code
    fr0=np.loadtxt('soldm_real.dat');
    fi0=np.loadtxt('soldm_imag.dat');

    fr1=np.real(f);
    fi1=np.imag(f);

    print "|fqk - fpy| = ",np.sqrt(np.sum(np.abs(fr0-fr1)**2))


    aps = {'params': { 'axes.labelsize':  10,
                       'text.fontsize':   10,
                       'legend.fontsize': 10,
                       'xtick.labelsize': 8,
                       'ytick.labelsize': 8,
                       'text.usetex': True,
                       'figure.figsize': (4,7)},
                 'axes': [0.125,0.2,0.95-0.125,0.95-0.2]}
    
    rcParams.update(aps['params'])    
    rcParams['text.dvipnghack']=True
    rcParams['font.family']='serif'
    rcParams['font.serif']='Computer Modern Roman'

    clf();
    subplot(2,1,1);
    plot(t,fr0)
    plot(t,fr1,'--');
    legend(["qk 00","qk 11","qk 01","qk 10","py 00","py 11","py 01","py 10"])
    ylabel("real")
    subplot(2,1,2);
    plot(t,fi0);
    plot(t,fi1,'--');
    xlabel("time")
    ylabel("imaginary")
    legend(["qk 00","qk 11","qk 01","qk 10","py 00","py 11","py 01","py 10"])
    
    savefig("test2x2.pdf")

